#include <thread>
#include <random>
#include <unordered_set>

#include "kvalloc.h"

#include "gtest/gtest.h"

using namespace kvalloc;

class multi_thread_test : public testing::Test {
protected:
    KVAlloc * allocator;

public:
    virtual void SetUp() {
        if(file_exist("/home/lyp/dnskv/build/test.db"))
            remove("/home/lyp/dnskv/build/test.db");
        if(file_exist("/mnt/pmem/testmeta.pool"))
            remove("/mnt/pmem/testmeta.pool");
        allocator = new KVAlloc("/home/lyp/dnskv/build/test.db", "/mnt/pmem/testmeta.pool", 4096);
    }

    virtual void TearDown() {
        delete allocator;
    }
};

TEST_F(multi_thread_test, alloc) {
    const int thread_num  = 4;
    uint64_t slabid_array[thread_num];

    const int ALLOC_SIZE = 120;
    auto iter = SIZECLASS.lower_bound(ALLOC_SIZE);
    int TEST_NUM = SLAB_SIZE / iter->first + 1;
    
    auto func = [&](int thread_id) {
        for(int k = 0; k < 5; k++) {
            FileAddr tmp[TEST_NUM];
            for(int i = 0; i < TEST_NUM; i++) {
                tmp[i] = allocator->Alloc(ALLOC_SIZE);
            }
            slabid_array[thread_id] = tmp[0].slabid;
            for(int i = 0; i < TEST_NUM; i++) {
                allocator->Free(tmp[i]);
            }
        }
    };

    std::thread ths[thread_num];
    for(int i = 0; i < thread_num; i++) {
        ths[i] = std::move(std::thread(func, i));
    }
    for(int i = 0; i < thread_num; i++) {
        ths[i].join();
    }
}

TEST_F(multi_thread_test, ranalloc) {
    const int total_num   = 4096;
    const int thread_num  = 4;
    FileAddr tmp[total_num];
    
    auto func = [&](int thread_id) {
        int perthread = total_num / thread_num;
        int offset = perthread * thread_id;
        std::default_random_engine gen;
        std::normal_distribution<double> dist(MAX_SIZECLASS / 2, MAX_SIZECLASS / 8);
        for(int i = 0; i < perthread; i++) {
            uint64_t alloc_size = dist(gen);
            if(alloc_size > MAX_SIZECLASS) alloc_size = MAX_SIZECLASS;
            if(alloc_size < MIN_SIZECLASS) alloc_size = MIN_SIZECLASS;
            tmp[offset + i] = allocator->Alloc(alloc_size);
        }
    };

    std::thread ths[thread_num];
    for(int i = 0; i < thread_num; i++) {
        ths[i] = std::move(std::thread(func, i));
    }
    for(int i = 0; i < thread_num; i++) {
        ths[i].join();
    }

    std::unordered_set<uint64_t> addr;
    for(int i = 0; i < (total_num / thread_num) * thread_num; i++) {
        ASSERT_TRUE(addr.insert(tmp[i].slabid * SLAB_SIZE + tmp[i].slaboff).second);
        allocator->Free(tmp[i]);
    }

    for(int i = 0; i < allocator->max_snum_; i++) {
        ASSERT_TRUE(allocator->at(i)->Empty());
    }
}

TEST_F(multi_thread_test, hybrid) {
    const int initial_num = 1024;
    FileAddr tmp[initial_num];
    std::default_random_engine gen;
    std::normal_distribution<double> dist(MAX_SIZECLASS / 3, MAX_SIZECLASS / 12);

    for(int i = 0; i < initial_num; i++) {
        uint64_t alloc_size = dist(gen);
        if(alloc_size > MAX_SIZECLASS) alloc_size = MAX_SIZECLASS;
        if(alloc_size < MIN_SIZECLASS) alloc_size = MIN_SIZECLASS;
        tmp[i] = allocator->Alloc(alloc_size);
    }

    const int thread_num  = 4;
    auto func = [&](int alloc_size) {
        auto iter = SIZECLASS.lower_bound(alloc_size);
        const int pcount = SLAB_SIZE / iter->first;
        FileAddr prvaite[pcount];
        for(int i = 0; i < pcount; i++) {
            prvaite[i] = allocator->Alloc(alloc_size);
        }
        for(int i = 0; i < pcount; i++) {
            allocator->Free(prvaite[i]);
        }
    };
    std::thread ths[thread_num];
    for(int i = 0; i < thread_num; i++) {
        ths[i] = std::move(std::thread(func, i * 47));
    }
    for(int i = 0; i < thread_num; i++) {
        ths[i].join();
    }

    for(int i = 0; i < initial_num; i++) {
        allocator->Free(tmp[i]);
    }
    
    for(int i = 0; i < allocator->max_snum_; i++) {
        ASSERT_TRUE(allocator->at(i)->Empty());
    }
}